# kim

import os
import argparse
from locust import HttpUser, task, between, events
from locust.env import Environment
from locust.argument_parser import LocustArgumentParser

# 假设以下模块是你的工具模块
from config.all_paths import get_root_path
from config.global_variable import ENV_DICT
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from utils.performance_tool.config_loader import load_config
from utils.performance_tool.load_shapes import SHAPE_MAP
from utils.request_tool.send_request import RequestSender
from utils.time_tool.time_control import parse_duration


# 添加自定义命令行参数
def add_custom_arguments(parser):
    parser.add_argument("--env", type=str, default="test_wt", help="Environment name")
    parser.add_argument("--role", type=str, default="doctor", help="Role name")
    parser.add_argument("--profile_name", type=str, required=True, help="Profile name")


# 获取自定义参数
def get_custom_params(environment):
    return {
        "env": environment.parsed_options.env,
        "role": environment.parsed_options.role,
        "profile_name": environment.parsed_options.profile_name,
    }


# 初始化测试
@events.test_start.add_listener
def init_test(environment, **kwargs):
    custom_params = get_custom_params(environment)
    env = custom_params["env"]
    role = custom_params["role"]
    profile_name = custom_params["profile_name"]

    # 加载配置
    config = load_config(env=env, role=role)
    profile_config = config["profiles"].get(profile_name, {})

    # 加载Excel测试数据
    excel_name = 'data/case_excel/wutong/uat/医生账号相关接口.xlsx'
    test_data_path = os.path.join(get_root_path(), excel_name)
    global testData
    testData = ExcelUtil(test_data_path).read_excel()

    print(f"\n=== 启动 {profile_name} 测试 ===")
    print(f"目标环境: {env}")
    print(f"角色: {role}")
    print(f"基准URL: {config['env'].get('url')}")
    print(f"测试数据加载路径: {test_data_path}")


# 测试收尾
@events.test_stop.add_listener
def finalize_test(environment, **kwargs):
    print("\n=== 测试结果摘要 ===")
    print(f"总请求数: {environment.stats.total.num_requests}")
    print(f"失败率: {environment.stats.total.fail_ratio:.2%}")


class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    abstract = True

    def on_start(self):
        custom_params = get_custom_params(self.environment)
        env = custom_params["env"]
        role = custom_params["role"]
        profile_name = custom_params["profile_name"]

        config = load_config(env=env, role=role)
        profile_config = config["profiles"].get(profile_name, {})

        self.s = Authentication(
            ENV_DICT.get("username"),
            ENV_DICT.get("password"),
            ENV_DICT.get("url"),
        ).auth_wutong()


class BaselineUser(BaseUser):
    """恒定负载"""

    @task
    def health_check(self):
        RequestSender().send_request(self.s, testData[0])


class StressUser(BaseUser):
    """峰值压力"""

    @task(5)
    def create_data(self):
        RequestSender().send_request(self.s, testData[0])


class ConcurrencyUser(BaseUser):
    """瞬时并发"""

    @task
    def update_resource(self):
        custom_params = get_custom_params(self.environment)
        profile_name = custom_params["profile_name"]

        profile_config = load_config(env=custom_params["env"], role=custom_params["role"])[
            "profiles"
        ].get(profile_name, {})

        resource_id = profile_config.get("concurrency_key")
        self.client.put(f"/resources/{resource_id}")


class WaveUser(BaseUser):
    """波浪型负载"""

    @task
    def wave_task(self):
        self.client.get("/api/wave-test")


# 动态测试类
class DynamicTestClass(eval(profile_config.get("user_class", "BaselineUser"))):
    pass


# 负载形状配置
class CustomLoadShape(SHAPE_MAP[profile_config.get("load_shape", "constant")]):
    custom_params = get_custom_params(self.environment)
    profile_name = custom_params["profile_name"]
    profile_config = load_config(env=custom_params["env"], role=custom_params["role"])[
        "profiles"
    ].get(profile_name, {})

    shape_users = profile_config.get("users", 1)
    spawn_rate = profile_config.get("spawn_rate", 1)
    shape_duration = parse_duration(profile_config.get("duration", "1m"))
    shape_stages = profile_config.get("stages", [])


# 暴露给Locust的配置
class TestScenario(DynamicTestClass):
    pass


def get_shape_class():
    custom_params = get_custom_params(self.environment)
    profile_name = custom_params["profile_name"]
    profile_config = load_config(env=custom_params["env"], role=custom_params["role"])[
        "profiles"
    ].get(profile_name, {})

    return SHAPE_MAP[profile_config.get("load_shape", "constant")]


# 扩展 Locust 的命令行参数解析器
if __name__ == "__main__":
    from locust import main

    parser = LocustArgumentParser()
    add_custom_arguments(parser)
    options = parser.parse_args()
    main.main(options)